#ifndef EMP_PRP_H__
#define EMP_PRP_H__
#include "emp-tool/utils/block.h"
#include "emp-tool/utils/constants.h"
#include "emp-tool/utils/aes.h"

namespace emp {

class PRP { public:
	AES_KEY aes;

	PRP(const char * seed = fix_key) {
		aes_set_key(seed);
	}

	PRP(const block& seed): PRP((const char *)&seed) {
	}

	void aes_set_key(const char * key) {
		aes_set_key(_mm_loadu_si128((block*)key));
	}

	void aes_set_key(const block& v) {
		AES_set_encrypt_key(v, &aes);
	}

	void permute_block(block *data, int nblocks) {
		for(int i = 0; i < nblocks/AES_BATCH_SIZE; ++i) {
			AES_ecb_encrypt_blks<AES_BATCH_SIZE>(data + i*AES_BATCH_SIZE, &aes);
		}
		int remain = nblocks % AES_BATCH_SIZE;
		AES_ecb_encrypt_blks(data + nblocks - remain, remain, &aes);
	}


       block H(block in, uint64_t id) {
		in = sigma(in);
		__m128i k_128 = _mm_loadl_epi64( (__m128i const *) (&id));
		in = xorBlocks(in, k_128);
		block t = in;
		permute_block(&t, 1);
		in =  xorBlocks(in, t);
		return in;	
	}

	template<int n>
		void H(block out[n], block in[n], uint64_t id) {
			block scratch[n];
			for(int i = 0; i < n; ++i) {
				out[i] = scratch[i] = xorBlocks(sigma(in[i]), _mm_loadl_epi64( (__m128i const *) (&id)));
				++id;
			}
			permute_block(scratch, n);
			xorBlocks_arr(out, scratch, out, n);
		}

	void Hn(block*out, block* in, uint64_t id, int length, block * scratch = nullptr) {
		bool del = false;
		if(scratch == nullptr) {
			del = true;
			scratch = new block[length];
		}
		for(int i = 0; i < length; ++i){
			out[i] = scratch[i] = xorBlocks(sigma(in[i]), _mm_loadl_epi64( (__m128i const *) (&id)));
			++id;
		}
		permute_block(scratch, length);
		xorBlocks_arr(out, scratch, out, length);
		if(del) {
			delete[] scratch;
			scratch = nullptr;
		}
	}
};
}
#endif// PRP_H__
